package com.kanq.memory;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.*;
import org.springframework.ai.content.Media;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.util.MimeType;

import java.io.IOException;
import java.net.URL;
import java.util.*;
import java.util.stream.Collectors;

public class RedisChatMemoryRepository implements ChatMemoryRepository {
    private final StringRedisTemplate stringRedisTemplate;  // 用于操作 Redis
    private final ObjectMapper objectMapper;    // 用于序列化和反序列化
    private final String PREFIX;       // 存储对话的 Redis Key 前缀
    private final String CONVERSATION_IDS_SET;  // 存储对话ID的 Redis Key

    public RedisChatMemoryRepository(StringRedisTemplate stringRedisTemplate, ObjectMapper objectMapper) {
        this(stringRedisTemplate, objectMapper, "chat:conversation:", "chat:all_conversation_ids");
    }

    public RedisChatMemoryRepository(StringRedisTemplate stringRedisTemplate, ObjectMapper objectMapper, String PREFIX) {
        this(stringRedisTemplate, objectMapper, PREFIX, "chat:all_conversation_ids");
    }

    public RedisChatMemoryRepository(StringRedisTemplate stringRedisTemplate, ObjectMapper objectMapper, String PREFIX, String CONVERSATION_IDS_SET) {
        this.stringRedisTemplate = stringRedisTemplate;
        this.objectMapper = objectMapper;
        this.PREFIX = PREFIX;
        this.CONVERSATION_IDS_SET = CONVERSATION_IDS_SET;
    }

    // 获取所有 conversationId（KEYS 命令匹配 chat:*）
    @Override
    public List<String> findConversationIds() {
        // 使用ZSet存储对话ID（更高效）
        // 获取对话ID集合（按时间倒序排序，越晚创建的对话ID排在前面）
        Set<String> conversationIds = stringRedisTemplate.opsForZSet().reverseRange(CONVERSATION_IDS_SET, 0, -1);

        if (conversationIds == null || conversationIds.isEmpty()) {
            return List.of();
        }

        return new ArrayList<>(conversationIds);
    }

    // 根据 conversationId 获取 Message 列表
    @Override
    public List<Message> findByConversationId(String conversationId) {
        // 参数验证
        if (conversationId == null || conversationId.isEmpty()) {
            throw new IllegalArgumentException("conversationId cannot be null or empty");
        }

        List<String> list = stringRedisTemplate.opsForList().range(PREFIX + conversationId, 0, -1);
        if (list == null || list.isEmpty()) {
            return List.of();
        }
        return list.stream()
                .map(json -> {
                    try {
                        // return objectMapper.convertValue(json, Message.class);   // 直接反序列化Message会报错
                        return deserializeMessage(json);    // 手动反序列化
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                })
                .collect(Collectors.toList());
    }

    // 保存整个 Message 列表到指定 conversationId
    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        // 参数验证
        if (conversationId == null || conversationId.isEmpty()) {
            throw new IllegalArgumentException("conversationId cannot be null or empty");
        }
        // 先清除原有的 conversation 数据
        stringRedisTemplate.delete(PREFIX + conversationId);

        if (messages == null || messages.isEmpty()) {
            return;
        }
        List<String> list = messages.stream()
                .map(message -> {
                    try {
                        return objectMapper.writeValueAsString(message);
                    } catch (JsonProcessingException e) {
                        throw new RuntimeException("Failed to serialize Message", e);
                    }
                })
                .collect(Collectors.toList());

        stringRedisTemplate.opsForList().rightPushAll(PREFIX + conversationId, list);
        // 更新对话ID集合
        stringRedisTemplate.opsForZSet().add(CONVERSATION_IDS_SET, conversationId, System.currentTimeMillis());
    }

    // 删除指定 conversationId 的数据
    @Override
    public void deleteByConversationId(String conversationId) {
        if (conversationId == null || conversationId.isEmpty()) {
            throw new IllegalArgumentException("conversationId cannot be null or empty");
        }

        stringRedisTemplate.delete(PREFIX + conversationId);
        stringRedisTemplate.opsForZSet().remove(CONVERSATION_IDS_SET, conversationId);
    }

    // 手动反序列化 Message
    public Message deserializeMessage(String json) throws IOException {
        // 解析 JSON 字符串为 JsonNode
        JsonNode jsonNode = objectMapper.readTree(json);
        // 获取 messageType 字段值
        if (!jsonNode.has("messageType")) {
            throw new IllegalArgumentException("Missing or invalid messageType field");
        }
        String messageType = jsonNode.get("messageType").asText();

        // 获取 text 字段值
        String text = jsonNode.has("text") ? jsonNode.get("text").asText() : "";

        // 获取 metadata 字段值
        Map<String, Object> metadata = getMetadata(jsonNode);

        // 获取 media 字段值
        List<Media> mediaList = getMediaList(jsonNode);

        return switch (MessageType.valueOf(messageType)) {
            case SYSTEM -> new SystemMessage(text);
            case USER -> UserMessage.builder()
                    .text(text)
                    .media(mediaList)
                    .metadata(metadata)
                    .build();
            case ASSISTANT -> {
                List<AssistantMessage.ToolCall> toolCalls = getToolCalls(jsonNode);
                yield new AssistantMessage(text, metadata, toolCalls, mediaList);
            }
            default -> throw new IllegalArgumentException("Unknown message type: " + messageType);
        };
    }

    private Media deserializeMedia(ObjectMapper mapper, JsonNode mediaNode) throws IOException {
        Media.Builder builder = Media.builder();

        // Handle MIME type
        if (mediaNode.has("mimeType")) {
            JsonNode mimeNode = mediaNode.get("mimeType");
            String type = mimeNode.get("type").asText();
            String subtype = mimeNode.get("subtype").asText();
            builder.mimeType(new MimeType(type, subtype));
        }

        // Handle data - could be either URL string or byte array
        if (mediaNode.has("data")) {
            String data = mediaNode.get("data").asText();
            if (data.startsWith("http://") || data.startsWith("https://")) {
                builder.data(new URL(data));
            } else {
                // Assume it's base64 encoded binary data
                byte[] bytes = Base64.getDecoder().decode(data);
                builder.data(bytes);
            }
        }

        // Handle dataAsByteArray if present (overrides data if both exist)
        if (mediaNode.has("dataAsByteArray")) {
            byte[] bytes = Base64.getDecoder().decode(mediaNode.get("dataAsByteArray").asText());
            builder.data(bytes);
        }

        // Handle optional fields
        if (mediaNode.has("id")) {
            builder.id(mediaNode.get("id").asText());
        }

        if (mediaNode.has("name")) {
            builder.name(mediaNode.get("name").asText());
        }

        return builder.build();
    }

    private Map<String, Object> getMetadata(JsonNode jsonNode) {
        if (jsonNode.has("metadata")) {
            return objectMapper.convertValue(jsonNode.get("metadata"), new TypeReference<>() {
            });
        }
        return new HashMap<>();
    }

    private List<Media> getMediaList(JsonNode jsonNode) throws IOException {
        List<Media> mediaList = new ArrayList<>();
        if (jsonNode.has("media")) {
            for (JsonNode mediaNode : jsonNode.get("media")) {
                mediaList.add(deserializeMedia(objectMapper, mediaNode));
            }
        }
        return mediaList;
    }

    private List<AssistantMessage.ToolCall> getToolCalls(JsonNode jsonNode) {
        if (jsonNode.has("toolCalls")) {
            return objectMapper.convertValue(jsonNode.get("toolCalls"), new TypeReference<>() {
            });
        }
        return Collections.emptyList();
    }
}